package cn.edu.sjtu.stap.cg.builder;

import java.util.*;

import org.aspectj.ajdt.internal.compiler.ast.*;
import org.aspectj.org.eclipse.jdt.internal.compiler.*;
import org.aspectj.org.eclipse.jdt.internal.compiler.ast.*;
import org.aspectj.org.eclipse.jdt.internal.compiler.lookup.*;

import cn.edu.sjtu.stap.cg.*;
import cn.edu.sjtu.stap.hg.*;
import cn.edu.sjtu.stap.util.*;

/**
 * CG builder for methods, constructors, and advices, etc.
 * 
 * @author William
 * 
 */
public class PartialCgBuilder extends ASTVisitor {
    private Map<AbstractMethodDeclaration, CallGraph> sourceMap = null;

    private Map<MethodBinding, CallGraph> binaryMap = null;

    private Set<AbstractMethodDeclaration> programEntries = null;

    private Stack<CallGraph> lastCgs = null;

    private CallGraph currentCg = null;

    private HierarchyGraph hg = null;

    public PartialCgBuilder() {
        initialize();
    }

    public void setHierarchyGraph(HierarchyGraph hg) {
        if (hg == null) {
            throw new NullPointerException("Illegal hierarchy graph!");
        }
        this.hg = hg;
    }

    public void initialize() {
        sourceMap = new HashMap<AbstractMethodDeclaration, CallGraph>();
        binaryMap = new HashMap<MethodBinding, CallGraph>();
        programEntries = new HashSet<AbstractMethodDeclaration>();
        lastCgs = new Stack<CallGraph>();
        currentCg = null;
    }

    public Map<AbstractMethodDeclaration, CallGraph> getResult() {
        return sourceMap;
    }

    /**
     * Get the program entries of all the CompilationUnitDeclarations.
     * 
     * @return Set of entry methods.
     */
    public Set<AbstractMethodDeclaration> getProgramEntries() {
        return programEntries;
    }

    private CallGraph findMethod(ASTNode node) {
        if (sourceMap.containsKey(node)) {
            return sourceMap.get(node);
        }
        if (node instanceof AbstractMethodDeclaration) {
            MethodBinding m = ((AbstractMethodDeclaration) node).binding;
            if (binaryMap.containsKey(m)) {
                return binaryMap.get(m);
            }
        }

        return null;
    }

    private CallGraph findMethod(MethodBinding methodBinding) {
        if (binaryMap.containsKey(methodBinding)) {
            return binaryMap.get(methodBinding);
        }

        return null;
    }

    private CallGraph addSourceMethod(AbstractMethodDeclaration node) {
        CallGraph g = sourceMap.get(node);
        if (null != g) {
            return g;
        }

        g = new CallGraph(node);
        sourceMap.put(node, g);

        return g;
    }

    private CallGraph addBinaryMethod(MethodBinding methodBinding) {
        CallGraph g = binaryMap.get(methodBinding);
        if (null != g) {
            return g;
        }

        CgBinaryMethod binaryMethod = new CgBinaryMethod();
        binaryMethod.setBinaryMethod(methodBinding);
        g = new CallGraph(binaryMethod);
        binaryMap.put(methodBinding, g);

        return g;
    }

    private void processMethodCall(ASTNode call, SourceTypeBinding receiverType, MethodBinding methodBinding) {
        AbstractMethodDeclaration m = methodBinding.sourceMethod();

        if (null == m) {
            CgMethod caller = currentCg.getRoot();
            addBinaryCall(call, caller, methodBinding);
        } else {
            Set<AbstractMethodDeclaration> callees = hg.getCallees(receiverType, m);
            CgMethod caller = currentCg.getRoot();
            Iterator<AbstractMethodDeclaration> i = callees.iterator();
            while (i.hasNext()) {
                AbstractMethodDeclaration callee = i.next();
                addSourceCall(call, caller, callee);
            }
        }
    }

    private void addBinaryCall(ASTNode node, CgMethod caller, MethodBinding methodBinding) {
        CallGraph g = findMethod(methodBinding);
        if (null == g) {
            g = addBinaryMethod(methodBinding);
        }

        if (!currentCg.contains(caller, g.getRoot())) {
            addCall(caller, g.getRoot());
        }
    }

    private void addSourceCall(ASTNode node, CgMethod caller, AbstractMethodDeclaration callee) {
        CallGraph g = findMethod(callee);
        if (null == g) {
            g = addSourceMethod(callee);
        }

        if (!currentCg.contains(caller, g.getRoot())) {
            addCall((CgMethod) caller, g.getRoot());
        }
    }

    private void addCall(CgMethod caller, CgMethod callee) {
        AbstractMethodDeclaration callerMethod = (AbstractMethodDeclaration) caller.getASTNode();
        AbstractMethodDeclaration calleeMethod = (AbstractMethodDeclaration) callee.getASTNode();
        if (currentCg.find(callerMethod, calleeMethod) != null) {
            return;
        }

        CgCallSite cs = new CgCallSite(caller, callee);
        currentCg.add(cs);

        CgCall c1 = new CgCall();
        c1.setSourceNode(caller);
        c1.setDestinationNode(cs);
        currentCg.add(c1);

        CgCall c2 = new CgCall();
        c2.setSourceNode(cs);
        c2.setDestinationNode(callee);
        currentCg.add(c2);
    }

    public boolean visit(ConstructorDeclaration constructorDeclaration, ClassScope scope) {
        lastCgs.push(currentCg);
        currentCg = findMethod(constructorDeclaration);
        if (null == currentCg) {
            currentCg = addSourceMethod(constructorDeclaration);
        }

        return true;
    }

    public void endVisit(ConstructorDeclaration constructorDeclaration, ClassScope scope) {
        currentCg = lastCgs.pop();
    }

    public boolean visit(MethodDeclaration methodDeclaration, ClassScope scope) {
        if (methodDeclaration instanceof PointcutDeclaration) {
            return false;
        }

        lastCgs.push(currentCg);
        currentCg = findMethod(methodDeclaration);
        if (null == currentCg) {
            currentCg = addSourceMethod(methodDeclaration);
        }

        // If this method is main method (or program entry), store it.
        MethodBinding binding = methodDeclaration.binding;
        if (null != binding) {
            if (binding.isMain()) {
                programEntries.add(methodDeclaration);
            }
        }

        return true;
    }

    public void endVisit(MethodDeclaration methodDeclaration, ClassScope scope) {
        if (methodDeclaration instanceof PointcutDeclaration) {
            return;
        }

        currentCg = lastCgs.pop();
    }

    public boolean visit(ExplicitConstructorCall explicitConstructor, BlockScope scope) {
        if (explicitConstructor.binding.declaringClass instanceof SourceTypeBinding) {
            processMethodCall(explicitConstructor, (SourceTypeBinding) explicitConstructor.binding.declaringClass,
                    explicitConstructor.binding);
        } else {
            // System.out.println("Not SourceTypeBinding: " +
            // explicitConstructor);
        }
        CallMap.getInstance()
                .addCall((AbstractMethodDeclaration) currentCg.getRoot().getASTNode(), explicitConstructor);
        return false;
    }

    public void endVisit(ExplicitConstructorCall explicitConstructor, BlockScope scope) {
        // Do nothing by default.
    }

    public boolean visit(MessageSend messageSend, BlockScope scope) {
        if (messageSend.receiver.resolvedType instanceof SourceTypeBinding) {
            processMethodCall(messageSend, (SourceTypeBinding) messageSend.receiver.resolvedType, messageSend.binding);
        } else {
            // System.err.println("Not SourceTypeBinding: " + messageSend);
        }
        CallMap.getInstance().addCall((AbstractMethodDeclaration) currentCg.getRoot().getASTNode(), messageSend);
        return true;
    }

    public void endVisit(MessageSend messageSend, BlockScope scope) {
        // Do nothing by default.
    }

    public boolean visit(AllocationExpression allocationExpression, BlockScope scope) {
        if (allocationExpression.resolvedType instanceof SourceTypeBinding) {
            processMethodCall(allocationExpression, (SourceTypeBinding) allocationExpression.resolvedType,
                    allocationExpression.binding);
        } else {
            // System.err.println("Not SourceTypeBinding: " +
            // allocationExpression);
        }
        CallMap.getInstance().addCall((AbstractMethodDeclaration) currentCg.getRoot().getASTNode(),
                allocationExpression);
        return true;
    }

    public void endVisit(AllocationExpression allocationExpression, BlockScope scope) {
        // Do nothing by default.
    }

    public boolean visit(FieldDeclaration fieldDeclaration, MethodScope scope) {
        return false;
    }
}
